/*
*    l3qd - Light, light, lightweight queue
*    Copyright (C) 2024  Marcus Pedersén marcus@marcux.org
*
*    This program is free software: you can redistribute it and/or modify
*    it under the terms of the GNU General Public License as published by
*    the Free Software Foundation, either version 3 of the License, or
*    (at your option) any later version.
*
*    This program is distributed in the hope that it will be useful,
*    but WITHOUT ANY WARRANTY; without even the implied warranty of
*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*    GNU General Public License for more details.
*
*    You should have received a copy of the GNU General Public License
*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package add

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/spf13/cobra"
	"notabug.org/marcux/l3q/internal/client"
	"notabug.org/marcux/l3q/internal/common"
	"os"
	"os/user"
	"slices"
	"strconv"
	"strings"
)

var addParaCmdFlagVerbose, addParaCmdFlagJobid bool
var addParaCmdFlagJobname, addParaCmdFlagDepend string
var addParaCmdFlagCores, addParaCmdFlagNodes int
var paraCoresArgInFile, paraDependArgInFile, paraJobnameArgInFile, paraNodesArgInFile bool

var addParaCmd = &cobra.Command{
	Use:   "para TASKFILE",
	Args:  cobra.ExactArgs(1),
	Short: "Adds a parallel job of single threaded tasks to l3q.",
	Long: `Description:
  Adds a new parallel job of tasks to l3q. 
  Number of tasks executed in parallel depends on number of cores and hosts specified.

  Positional argument TASKFILE is required.
  TASKFILE is path to file containing all tasks that should be run in parallel.
  One task per line, empty lines and lines starting with # are ignored. 
  Lines containing tasks must have one, two or three columns.
  If line only contains one column, that column specifies a path to an 
  executable program, no arguments can be used for this program and
  users home directory will be used as working directory.
  If line contains two columns, first column specifies a path to an 
  executable program, no arguments can be used for this program, 
  and the second column contains working directory but if first
  column is an integer this is interpretated as the number of cores and
  the second column is the program and users home directory will be 
  used as working directory.
  If line contains three columns, the first column must be an
  positive integer specifying number of cores that this specific task require.
  Column two contains the executable program and column three contains the
  working directory for the program.
  Lines that does not contain the cores column will have the number of cores
  set to 1 (one).
  Columns are separated by whitespace.

  --cores and --nodes are required arguments, either on commandline
  or as special syntax line in TASKFILE.
  If TASKFILE contains a line with the special syntax specifying required 
  options then no options are required on the command line.
  If TASKFILE contains the special syntax line and argument is 
  given on commandline this options will be used and the same options
  in special syntax line will be ignored.
  If TASKFILE contains multiple special syntax lines the last line will be used.

  Special syntax in TASKFILE:
  #--add-para --cores 5 --nodes 8`,
	Run: func(cmd *cobra.Command, args []string) {
		job := common.Job{
			Type: common.JobTypeParallel,
		}

		if addParaCmdFlagCores < 0 {
			fmt.Fprintf(os.Stderr, "Argument --cores is required to be a positive integer.\nTry -h or --help for help.\n")
			os.Exit(1)
		} else if addParaCmdFlagCores > 0 {
			job.CoresAlloc = addParaCmdFlagCores
		}

		if addParaCmdFlagNodes < 0 {
			fmt.Fprintf(os.Stderr, "Argument --nodesd is required to be a positive integer.\nTry -h or --help for help.\n")
			os.Exit(1)
		} else if addParaCmdFlagNodes > 0 {
			job.HostsAlloc = addParaCmdFlagNodes
		}

		if len(addParaCmdFlagDepend) > 0 {
			depStrs := strings.Split(strings.TrimSpace(addParaCmdFlagDepend), ",")

			error := false
			for _, d := range depStrs {
				if _, err := strconv.Atoi(d); err != nil {
					fmt.Fprintf(os.Stderr, "Failed to parse: %s as integer.\n", d)
					error = true
				}
			}

			if error {
				fmt.Fprintf(os.Stderr, "Argument --depend is required to be a comma separated list of integers\n")
				fmt.Fprintf(os.Stderr, "Try -h or --help for help.\n")

				os.Exit(1)
			}

			job.Depend = strings.TrimSpace(addParaCmdFlagDepend)
		}

		if len(addParaCmdFlagJobname) > 0 {
			job.Name = addParaCmdFlagJobname
		}

		usr, err := user.Current()
		if err != nil {
			if addParaCmdFlagVerbose {
				fmt.Fprintf(os.Stderr, "Failed to get current user.\n%s\n", err)
			} else {
				fmt.Fprintf(os.Stderr, "Failed to get current user.\n")
			}

			os.Exit(1)
		}

		job.User = usr.Username

		grp, err := user.LookupGroupId(usr.Gid)
		if err != nil {
			if addSeqCmdFlagVerbose {
				fmt.Fprintf(os.Stderr, "Failed to lookup current users main group.\n%s\n", err)
			} else {
				fmt.Fprintf(os.Stderr, "Failed to lookup current users masin group.\n")
			}

			os.Exit(1)
		}

		job.Group = grp.Name

		err = paraParseTaskFile(&job, args[0])
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to parse TASKFILE: \n%s\n", err)
			os.Exit(1)
		}

		if paraCoresArgInFile {
			if addParaCmdFlagCores > 0 {
				job.CoresAlloc = addParaCmdFlagCores
			}
		}

		if paraNodesArgInFile {
			if addParaCmdFlagNodes > 0 {
				job.HostsAlloc = addParaCmdFlagNodes
			}
		}

		if paraDependArgInFile {
			if len(addParaCmdFlagDepend) > 0 {
				job.Depend = strings.TrimSpace(addParaCmdFlagDepend)
			}
		}

		if paraJobnameArgInFile {
			if len(addParaCmdFlagJobname) > 0 {
				job.Name = addParaCmdFlagJobname
			}
		}

		if len(job.Name) == 0 {
			job.Name = client.GenerateJobName()
		}

		if job.CoresAlloc == 0 && job.HostsAlloc == 0 {
			fmt.Fprintf(os.Stderr, "Both --cores and --nodes flags are required.\nTry -h or --help for help.\n")
			os.Exit(1)
		} else if job.CoresAlloc == 0 {
			fmt.Fprintf(os.Stderr, "Flag --cores is required.\nTry -h or --help for help.\n")
			os.Exit(1)
		} else if job.HostsAlloc == 0 {
			fmt.Fprintf(os.Stderr, "Flag --nodes is required.\nTry -h or --help for help.\n")
			os.Exit(1)
		}

		taskCores := []int{}
		jobCoresLarger := true

		for _, t := range job.Tasks {
			if job.CoresAlloc < t.Cores {
				fmt.Fprintf(os.Stderr, "TASKFILE: %s contains tasks that has more cores specified then specified in job: %d\n",
					args[0], job.CoresAlloc)
				fmt.Fprintf(os.Stderr, "Number of cores specified in job is required to be at least as large as the largest\n")
				fmt.Fprintf(os.Stderr, "number for cores specified for the largest task in TASKFILE.\n")
				os.Exit(1)
			} else if job.CoresAlloc == t.Cores {
				jobCoresLarger = false
			}

			if !slices.Contains(taskCores, t.Cores) {
				taskCores = append(taskCores, t.Cores)
			}
		}

		slices.Sort(taskCores)

		if jobCoresLarger && len(taskCores) >= 2 {
			if job.CoresAlloc < taskCores[0]+taskCores[1] {
				fmt.Fprintf(os.Stderr, "Number of cores specified for job: %d is larger then number of cores specified for tasks\n",
					job.CoresAlloc)
				fmt.Fprintf(os.Stderr, "but not large enough to fit more than one task per node, unused resources will be allocated.\n")
				os.Exit(1)
			}
		} else if jobCoresLarger && len(taskCores) == 1 {
			if job.CoresAlloc%taskCores[0] != 0 && len(job.Tasks) > 1 {
				fmt.Fprintf(os.Stderr, "Number of cores specified for job: %d is larger then number of cores specified for tasks: %d\n",
					job.CoresAlloc, taskCores[0])
				fmt.Fprintf(os.Stderr, "but not large enough to fit the full allocation with tasks, unused resources will be allocated.\n")
				os.Exit(1)
			} else if len(job.Tasks) == 1 {
				fmt.Fprintf(os.Stderr, "Number of cores specified for job: %d is larger then number of cores specified for task: %d\n",
					job.CoresAlloc, taskCores[0])
				fmt.Fprintf(os.Stderr, "Unused resources will be allocated.\n")
				os.Exit(1)
			}
		}

		id := client.UrlPostJob("/job", addParaCmdFlagVerbose, job)

		if addParaCmdFlagJobid {
			fmt.Printf("%d\n", id)
		} else {
			fmt.Printf("Job with name: %s and Id: %d, sussessfully added.\n", job.Name, id)
		}
	},
}

func paraParseFileArgs(job *common.Job, argLine string) error {
	if !strings.HasPrefix(argLine, "#--add-para ") {
		return fmt.Errorf("Line does not start with required: #--add-para")
	}

	argLine, _ = strings.CutPrefix(argLine, "#--add-para ")
	args := strings.Fields(argLine)

	errStr := ""
	nextArg := true
	arg := ""
	for _, a := range args {
		if nextArg {
			switch a {
			case "-c", "--cores":
				arg = "--cores"
				nextArg = false
			case "-n", "--nodes":
				arg = "--nodes"
				nextArg = false
			case "-d", "--depend":
				arg = "--depend"
				nextArg = false
			case "-j", "--jobname":
				arg = "--jobname"
				nextArg = false
			case "-i", "--jobid":
				addParaCmdFlagJobid = true
				nextArg = true
			default:
				return fmt.Errorf("Wrong argument, argument: %s, unknown", a)
			}
		} else {
			switch arg {
			case "--cores":
				c, err := strconv.Atoi(a)
				if err != nil {
					errStr = fmt.Sprintf("%s\nFailed to parse argument --cores: %s, %s",
						errStr, a, err)
				} else if c <= 0 {
					errStr = fmt.Sprintf("%s\nArgument --cores: %s, is not a valid positiv integer.",
						errStr, a)
				} else {
					job.CoresAlloc = c
					paraCoresArgInFile = true
				}
			case "--nodes":
				n, err := strconv.Atoi(a)
				if err != nil {
					errStr = fmt.Sprintf("%s\nFailed to parse argument --nodes: %s, %s",
						errStr, a, err)
				} else if n <= 0 {
					errStr = fmt.Sprintf("%s\nArgument --nodes: %s, is not a valid positiv integer.",
						errStr, a)
				} else {
					job.HostsAlloc = n
					paraNodesArgInFile = true
				}
			case "--depend":
				depStrs := strings.Split(strings.TrimSpace(a), ",")

				error := false
				for _, d := range depStrs {
					if _, err := strconv.Atoi(d); err != nil {
						errStr = fmt.Sprintf("%s\nFailed to parse depend value: %s as integer.",
							errStr, d)
						error = true
					}
				}

				if error {
					errStr = fmt.Sprintf("%s\nArgument --depend: %s, is required to be a comma separated list of positive integers.", errStr, a)
				} else {
					job.Depend = strings.TrimSpace(a)
					paraDependArgInFile = true
				}
			case "--jobname":
				job.Name = a
				paraJobnameArgInFile = true
			}

			nextArg = true
		}
	}

	if len(errStr) > 0 {
		return fmt.Errorf("%s", errStr)
	}

	return nil
}

func paraParseTaskFile(job *common.Job, filePath string) error {
	file, err := os.Open(filePath)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return fmt.Errorf("File: %s can not be found.", filePath)
		} else {
			return fmt.Errorf("Failed to open file: %s, %s", filePath, err)
		}
	}

	defer file.Close()

	scanner := bufio.NewScanner(file)

	i := 0
	seqno := 0
	argLine := ""
	lineErr := ""
	for scanner.Scan() {
		i++
		line := scanner.Text()

		if strings.HasPrefix(line, "#--add-para ") {
			argLine = line
		} else if strings.HasPrefix(line, "#") {
			// comment
			continue
		} else if len(strings.TrimSpace(line)) == 0 {
			// empty line
			continue
		} else {
			// parse task
			fields := strings.Fields(line)
			if len(fields) > 3 {
				lineErr = fmt.Sprintf("%sLine: %d Too many fields in line. Max fields: %d, fields in line: %d\n",
					lineErr, i, 3, len(fields))
				continue
			}

			task := common.Task{}
			seqno++
			task.Seqno = seqno
			task.Cores = 1

			if len(fields) == 1 {
				task.Command = fields[0]
				task.WorkDir = "$HOME"
			} else if len(fields) == 2 {
				c, err := strconv.Atoi(fields[0])
				if err == nil {
					if c < 1 {
						lineErr = fmt.Sprintf("%sLine: %d First field is required to be a positive integer: %s\n",
							lineErr, i, fields[0])
						continue
					}

					task.Cores = c
					task.Command = fields[1]
					task.WorkDir = "$HOME"
				} else {
					task.Command = fields[0]
					task.WorkDir = fields[1]
				}

			} else {
				c, err := strconv.Atoi(fields[0])
				if err != nil {
					lineErr = fmt.Sprintf("%sLine: %d Failed to parse first field as an integer: %s\n",
						lineErr, i, fields[0])
					continue
				}

				if c < 1 {
					lineErr = fmt.Sprintf("%sLine: %d First field is required to be a positive integer: %s\n",
						lineErr, i, fields[0])
					continue
				}

				task.Command = fields[1]
				task.WorkDir = fields[2]
				task.Cores = c
			}

			job.Tasks = append(job.Tasks, task)
		}
	}

	if len(lineErr) > 0 {
		return fmt.Errorf("Error in %s:\n%s", filePath, lineErr)
	}

	if len(argLine) > 0 {
		err := paraParseFileArgs(job, argLine)
		if err != nil {
			return fmt.Errorf("Error in %s:\nFailed to parse argument line:\n%s\n%s",
				filePath, argLine, err)
		}
	}

	return nil
}

func init() {
	addParaCmd.Flags().BoolVarP(&addParaCmdFlagVerbose, "verbose", "v", false, "Print verbose output")
	addParaCmd.Flags().BoolVarP(&addParaCmdFlagJobid, "jobid", "i", false, "Jobid of this job will be returned on success\nas the only output.\nOn failure the error message is printed.")
	addParaCmd.Flags().StringVarP(&addParaCmdFlagJobname, "jobname", "j", "", "Name of specified job")
	addParaCmd.Flags().StringVarP(&addParaCmdFlagDepend, "depend", "d", "",
		`Comma separated list of other jobid that this job depends on.
All jobs specified in this joblist have to terminate without 
errors before this job is started. If a depenedent job fails 
this job will also fail and will not be started.
Example: 1,5,9 (comma separated without whitespaces)`)
	addParaCmd.Flags().IntVarP(&addParaCmdFlagCores, "cores", "c", 0,
		`Number of cores to use on each node, --nodes is also required.`)
	addParaCmd.Flags().IntVarP(&addParaCmdFlagNodes, "nodes", "n", 0,
		`Number of nodes to use, l3q will choose which nodes to use.
--cores is also required.`)
}
